#ifndef SRC_HYPERCUBE_HYPERCUBE_HMM_ENGINE
#define SRC_HYPERCUBE_HYPERCUBE_HMM_ENGINE

#include "h_mat_mul.hpp"
#include "helib_atm_wrapper.hpp"

class hypercube_hmme;

template <> struct hmm_status<hypercube_hmme> : public hmm_status<shmm_engine> {
public:
  hmm_status<hypercube_hmme>() = delete;
  hmm_status<hypercube_hmme>(const hmm_status<shmm_engine> &status);

  bool encoding_reuse() const { return encoding_reuse_; }
  bool clean_left() const { return clean_left_; }
  bool clean_right() const { return clean_right_; }

  void set_encoding_reuse(bool val) { encoding_reuse_ = val; }
  void set_clean_left(bool val) { clean_left_ = val; }
  void set_clean_right(bool val) { clean_right_ = val; }

  operator hmm_t() const;
  operator hmm_mt_t() const;

private:
  bool encoding_reuse_ = true;
  bool clean_left_ = false;
  bool clean_right_ = false;
};

// * hypercube_hmme(aka. shhmm_engine)
//   * s: single-ciphertext
//   * h: hypercube packing
//   * hmm: homomorphic matrix multiplication
class hypercube_hmme : public shmm_engine {
public:
  /**
   * @brief Construct a new `hypercube_hmme` object.
   *        See `shmm_engine` in `h_mat_mul.hpp` for more detials
   */
  hypercube_hmme(const hmm_status<shmm_engine> &status);

  hypercube_hmme(const hypercube_hmme &hmme);

  ~hypercube_hmme() = default;

  static cube_signature get_hypercube(long m, long n, long k);

  const helib::EncryptedArray &get_ea(MATRIX_TYPE type) const {
    return *ea_[type];
  }

  double get_encoding_time() const {
    // if (status()->encoding_reuse()) {
    //   return 0.0;
    // }
    double time = 0.0;
    for (int i = 0; i < UNDEF_MATRIX_TYPE; i++) {
      if (metric_collector_[i]) {
        time += metric_collector_[i]->get_no_encoding_reuse();
      }
    }
    return time;
  }

  /**
   * @brief Using `helib::EncryptedArray` to activate this engine
   *
   * @param ea  class for data-movement operations on arrays of slots
   */
  void register_engine(const helib::EncryptedArray &ea,
                       const cube_signature &cube) override;

  /**
   * @brief Execute Ciphertext GEMMs.
   *        See `shmm_engine` in `h_mat_mul.hpp` for more detials
   */
  void cgemm(MATRIX_TYPE type, helib::Ctxt &AB, const helib::Ctxt &A,
             const helib::Ctxt &B) override;

  void rotate_align_A(MATRIX_TYPE type, int rot_pos_A, int rot_size_A,
                      helib::Ctxt &Ap, const helib::Ctxt &A);
  void rotate_align_B(MATRIX_TYPE type, int rot_pos_A, int rot_size_A,
                      helib::Ctxt &Bp, const helib::Ctxt &B);
  void shift_compute_A(MATRIX_TYPE type, int rot_pos_A,
                       std::vector<helib::Ctxt> &Ais, const helib::Ctxt &A);
  void shift_compute_B(MATRIX_TYPE type, int rot_pos_B,
                       std::vector<helib::Ctxt> &Bis, const helib::Ctxt &B);
  void rotate_align_A(MATRIX_TYPE type, helib::Ctxt &Ap, const helib::Ctxt &A);
  void rotate_align_B(MATRIX_TYPE type, helib::Ctxt &Bp, const helib::Ctxt &B);
  helib::Ctxt shift_compute_A(
      MATRIX_TYPE type,
      const std::shared_ptr<helib::GeneralAutomorphPrecon_FULL> &hoisted_A,
      int r);

  hmm_status<hypercube_hmme> *status() override {
    return static_cast<hmm_status<hypercube_hmme> *>(status_);
  }

  const hmm_status<hypercube_hmme> *status() const {
    return static_cast<hmm_status<hypercube_hmme> *>(status_);
  }
};

#endif /* SRC_HYPERCUBE_HYPERCUBE_HMM_ENGINE */
